package common

import (
	"regexp"
	"strings"
	"time"
)

// Response 模型返回结果统一
type Response struct {
	model            string
	think            string
	answer           string
	promptTokens     int
	completionTokens int
	totalTokens      int
	tools            []Tool
	startTime        time.Time
	endTime          time.Time
	firstTime        time.Time
}

// GetModelType 获取模型类型
func (r *Response) GetModelType() string {
	return r.model
}

// GetThink 获取思考过程
func (r *Response) GetThink() string {
	if r.think == "" && strings.Contains(r.answer, "<think>") {
		think, answer := r.rex(r.answer)
		r.answer = answer
		r.think = think
	}
	return strings.TrimSpace(r.think)
}

// GetAnswer 获取模型回复
func (r *Response) GetAnswer() string {
	if strings.Contains(r.answer, "<think>") {
		think, answer := r.rex(r.answer)
		r.answer = answer
		r.think = think
	}
	return strings.TrimSpace(r.answer)
}

// GetTools 获取工具调用
func (r *Response) GetTools() []Tool {
	return r.tools
}

// GetTokens 获取token统计数据
func (r *Response) GetTokens() map[string]int {
	return map[string]int{
		"promptTokens":     r.promptTokens,
		"completionTokens": r.completionTokens,
		"totalTokens":      r.totalTokens,
	}
}

// setModelType 设置模型类型
func (r *Response) setModelType(model string) {
	r.model = model
}

// setThink 设置模型是考内容并返回
func (r *Response) setThink(think string) string {
	r.think += think
	return r.think
}

// setAnswer 设置模型回答结果并返回
func (r *Response) setAnswer(answer string) string {
	r.answer += answer
	return r.answer
}

func (r *Response) Update(res ChatResponse) {
	r.setModelType(res.Model)

	// 提取思考和结果
	if res.Choices != nil && len(res.Choices) > 0 {
		choices := res.Choices[0]
		var think, answer string
		if choices.Delta != nil {
			think = choices.Delta.ReasoningContent
			answer = choices.Delta.Content
			// 工具调用
			for i, call := range choices.Delta.ToolCalls {
				if len(r.tools) <= i {
					r.tools = append(r.tools, call)
				} else {
					r.tools[i].Function.Name += call.Function.Name
					r.tools[i].Function.Arguments += call.Function.Arguments
				}
			}
		} else if choices.Message != nil {
			content := strings.TrimSpace(choices.Message.Content)
			think, answer = r.rex(content)
			// 工具调用
			r.tools = choices.Message.ToolCalls
		}
		r.setThink(think)
		r.setAnswer(answer)
	}

	if &res.Usage != nil {
		r.setTokens(res.Usage.PromptTokens, res.Usage.CompletionTokens, res.Usage.TotalTokens)
	}
}

// setTokens 设置token统计数据
func (r *Response) setTokens(promptTokens, completionTokens, totalTokens int) {
	r.promptTokens = promptTokens
	r.completionTokens = completionTokens
	r.totalTokens = totalTokens
}

func (r *Response) rex(content string) (string, string) {
	re := regexp.MustCompile(`(?ims)<think>(.*?)</think>`)
	think := re.FindString(content)
	answer := strings.TrimPrefix(content, think)
	return think, answer
}

// SetTime 设置时间
func (r *Response) SetTime(mode string) {
	if mode == "start" {
		r.startTime = time.Now()
	} else if mode == "end" {
		r.endTime = time.Now()
	} else if mode == "first" && r.firstTime.IsZero() {
		r.firstTime = time.Now()
	}
}

// GetTime 获取时间
func (r *Response) GetTime() map[string]time.Time {
	return map[string]time.Time{
		"start": r.startTime,
		"end":   r.endTime,
		"first": r.firstTime,
	}
}

// ChatResponse 响应结构体（根据实际 API 响应格式调整）
type ChatResponse struct {
	ID      string `json:"id"`
	Object  string `json:"object"`
	Created int64  `json:"created"`
	Model   string `json:"model"`
	Choices []struct {
		// 非流式输出 非深度思考
		Message *struct {
			Role      string `json:"role"`
			Content   string `json:"content"`
			ToolCalls []Tool `json:"tool_calls"`
		} `json:"message,omitempty"`
		// 流式输出，深度思考
		Delta *struct {
			Content          string `json:"content"`
			ReasoningContent string `json:"reasoning_content"`
			ToolCalls        []Tool `json:"tool_calls"`
		} `json:"delta,omitempty"`
		FinishReason string `json:"finish_reason"`
	} `json:"choices"`
	Usage struct {
		PromptTokens     int `json:"prompt_tokens"`
		CompletionTokens int `json:"completion_tokens"`
		TotalTokens      int `json:"total_tokens"`
	} `json:"usage"`
}
